from utils import print_timing, is_blankline

class Anno_Base(object):
    def __init__(self, annotfile, schema, tmpdir, table_dict):
        self.annotfile = annotfile
        self.schema = schema
        self.tmpdir = tmpdir
        self.table_dict = table_dict

    def discard_lines_before_annotate_header(self):
        import string
        af = open(self.annotfile)
        discard_lines=0
        discard_chars = string.whitespace + ','
        while True:
            line = af.readline()
            line_split = line.lstrip(discard_chars).rstrip(discard_chars).split(',')
            if(len(line_split) > 10):
                break
            if not line or discard_lines > 10000:
                import sys
                print "Couldn't find annotate header in the first %s lines of %s"%(discard_lines, self.annotfilename)
                print "Is it a correctly formatted annotation file?"
                sys.exit(1)
            discard_lines += 1
        af.close()
        return discard_lines

    def chromosome_sub(self, chrom):
        if chrom == "X":
            return 23
        elif chrom == "Y":
            return 24
        elif chrom == "XY":
            return 25
        elif chrom == "MT" or chrom == "Mt":
            return 26
        else:
            return int(chrom)

    def sort_anno_dict_list(self, x, y):
        return cmp(x['chromosome'], y['chromosome']) or cmp(x['location'], y['location']) or cmp(x['fid'], y['fid'])

    @print_timing
    def make_anno_table(self, session, alleled):
        import os
        print "making anno table..."
        anno_table = self.table_dict['anno_table']
        from dbutils import add_obj
        index_dict = self.data_indices()
        dirpath = os.path.dirname(os.path.abspath(__file__))
        annotablefile = os.path.abspath(os.path.join(self.tmpdir, self.annotablefile))
        self.write_anno_table_to_file(session, annotablefile, alleled)
        self.drop_anno_table_constraints_and_indexes(session)
        print "writing anno table data to database..."
        print "writing csv file %s to anno table..."%annotablefile
        session.execute("COMMIT; BEGIN; TRUNCATE %s.anno; COPY %s.anno FROM '%s' USING DELIMITERS ','; COMMIT;"%(self.schema, self.schema, annotablefile))
        self.restore_anno_table_constraints_and_indexes(session)
        print "finished writing anno table..."

class Anno_Affy6(Anno_Base):
    def __init__(self, annotfile, schema, tmpdir, table_dict):
        Anno_Base.__init__(self, annotfile, schema, tmpdir, table_dict)
        self.annotablefile = "anno_affy6_%s"%self.schema

    def data_indices(self):
        af = open(self.annotfile)
        discard_lines = self.discard_lines_before_annotate_header()
        for i in range(discard_lines):
            af.readline()
        headers = af.readline().split(',')
        for h, i in zip(headers, range(len(headers))):
            h = h.lstrip('"').rstrip('"')
            if "Allele A" == h:
                alleleA = i
            if "Allele B" == h:
                alleleB = i
            if  "Chromosome" == h:
                chromosome = i
            if "Physical Position" == h:
                location = i
            if "Probe Set ID" == h:
                probe_set_id = i
            if "dbSNP RS ID" == h:
                rsid = i
        af.close()
        return {'alleleA':alleleA, 'alleleB':alleleB, 'chromosome':chromosome, 'location':location, 'probe_set_id':probe_set_id, 'rsid':rsid}

    @print_timing
    def write_anno_table_to_file(self, session, annotablefile, alleled):
        print "writing csv file %s containing affy6 anno table data..."%annotablefile
        index_dict = self.data_indices()
        alleleA, alleleB, chromosome, location, probe_set_id, rsid = map(lambda x: index_dict[x], ["alleleA", "alleleB", "chromosome", "location", "probe_set_id", "rsid"])
        af = open(self.annotfile)
        atf = open(annotablefile, 'w')
        discard_lines = self.discard_lines_before_annotate_header();
        for i in range(discard_lines):
            af.readline()
        header = af.readline()
        i = 0
        insert_dict_list = []
        print "reading lines from calls file..."
        for line in af:
            i = i+1
            line = line.lstrip().rstrip().split(',')
            fidval = line[probe_set_id].lstrip('"').rstrip('"')
            rsidval = line[rsid].lstrip('"').rstrip('"')
            chromval = line[chromosome].lstrip('"').rstrip('"')
            locationval = line[location].lstrip('"').rstrip('"')
            alleleAval =  line[alleleA].lstrip('"').rstrip('"')
            alleleBval =  line[alleleB].lstrip('"').rstrip('"')
            if chromval != '---' and rsidval != '---':
                # make sure both chromosome and location values are integers or sorting won't work.
                d = {'fid': fidval, 'rsid' : rsidval, 'chromosome': self.chromosome_sub(chromval),
                     'location': int(locationval), 'allelea_id': alleleAval, 'alleleb_id': alleleBval}
                insert_dict_list.append(d)
        print "done reading lines from calls file, sorting anno table data before insert into database..."
        insert_dict_list.sort(self.sort_anno_dict_list)
        for i, d in enumerate(insert_dict_list):
            atf.write(','.join([str(i+1), d['fid'], d['rsid'], str(d['chromosome']), str(d['location']), d['allelea_id'], d['alleleb_id']])+'\n')
        print "finished writing anno table to file..."
        atf.close()

    @print_timing
    def drop_anno_table_constraints_and_indexes(self, session):
        print "dropping anno table constraints..."
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_fid_key', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_pkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_allelea_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_alleleb_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_chromosome_fkey', '%s');"%self.schema)
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_allelea_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_alleleb_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_chromosome;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_location;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_rsid;"%(self.schema, self.schema))
        session.commit()
        session.flush()

    @print_timing
    def restore_anno_table_constraints_and_indexes(self, session):
        print "restoring anno table constraints..."
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_fid_key UNIQUE (fid);"%(self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_pkey PRIMARY KEY (id);"%(self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_allelea_id_fkey FOREIGN KEY (allelea_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_alleleb_id_fkey FOREIGN KEY (alleleb_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_chromosome_fkey FOREIGN KEY (chromosome) REFERENCES %s.chromo(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_allelea_id ON %s.anno USING btree (allelea_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_alleleb_id ON %s.anno USING btree (alleleb_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_chromosome ON %s.anno USING btree (chromosome);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_location ON %s.anno USING btree (location);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_rsid ON %s.anno USING btree (rsid);"%(self.schema, self.schema))
        session.commit()
        session.flush()

# Starting with X, Y (values in the annotation file corresponding to the
# 'SNP' field in the Illumina annotation file)

# Let F: {A, C, G, T} -> {A, C, G, T} be defined by

# C -> G
# G -> C
# A -> T
# T -> A

# So F is the reverse complement. In the following Aforward, Atop are
# the A alleles for forward and top respectively, etc.

# istrand     sstrand     Aforward      Bforward       Atop         Btop
#  BOT         BOT           X             Y            F(X)         F(Y)
#  BOT         TOP         F(X)          F(Y)           F(X)         F(Y)
#  TOP         BOT         F(X)          F(Y)             X            Y
#  TOP         TOP           X             Y              X            Y

# Once I have the A and B alleles, I can just read off the integer call
# values (or whatever the correct term is) by doing

# AA -> 0
# AB -> 1
# BB -> 2

# Starting with X, Y
# Let F: {A, C, G, T} -> {A, C, G, T} be defined by

# C -> G
# G -> C
# A -> T
# T -> A

# istrand     sstrand     Aforward      Bforward       Atop         Btop
# BOT         BOT           X             Y            F(X)         F(Y)
# BOT         TOP         F(X)          F(Y)           F(X)         F(Y)
# TOP         BOT         F(X)          F(Y)             X            Y
# TOP         TOP           X             Y              X            Y

class Anno_Illumina(Anno_Base):
    def __init__(self, annotfile, schema, tmpdir, table_dict):
        Anno_Base.__init__(self, annotfile, schema, tmpdir, table_dict)
        self.annotablefile = "anno_illumina_%s"%self.schema

    def data_indices(self):
        af = open(self.annotfile)
        discard_lines = self.discard_lines_before_annotate_header()
        for i in range(discard_lines):
            af.readline()
        headers = af.readline().split(',')
        for h, i in zip(headers, range(len(headers))):
            h = h.lstrip('"').rstrip('"')
            if "SNP" == h:
                snp = i
            if "IlmnStrand" == h:
                istrand = i
            if "SourceStrand" == h:
                sstrand = i
            if  "Chr" == h:
                chromosome = i
            if "MapInfo" == h:
                location = i
            if "IlmnID" == h:
                probe_set_id = i
            if "Name" == h:
                rsid = i
        af.close()
        return {'snp':snp, 'istrand':istrand, 'sstrand':sstrand, 'chromosome':chromosome, 'location':location, 'probe_set_id':probe_set_id, 'rsid':rsid}

    def reverse_complement(self, alleleA, alleleB):
        import string
        table = string.maketrans('ACGT', 'TGCA')
        return alleleA.translate(table), alleleB.translate(table)

    def alleles_trans(self, allelea, alleleb, istrand, sstrand, alleletype):
        if alleletype == "forward":
            if istrand == 'BOT' and sstrand == 'BOT':
                return allelea, alleleb
            if istrand == 'BOT' and sstrand == 'TOP':
                return self.reverse_complement(allelea, alleleb)
            if istrand == 'TOP' and sstrand == 'BOT':
                return self.reverse_complement(allelea, alleleb)
            if istrand == 'TOP' and sstrand == 'TOP':
                return allelea, alleleb
        elif alleletype == "top":
            if istrand == 'BOT' and sstrand == 'BOT':
                return self.reverse_complement(allelea, alleleb)
            if istrand == 'BOT' and sstrand == 'TOP':
                return self.reverse_complement(allelea, alleleb)
            if istrand == 'TOP' and sstrand == 'BOT':
                return allelea, alleleb
            if istrand == 'TOP' and sstrand == 'TOP':
                return allelea, alleleb

    @print_timing
    def write_anno_table_to_file(self, session, annotablefile, alleled):
        import string
        print "writing csv file %s containing illumina anno table data..."%annotablefile
        index_dict = self.data_indices()
        snp, istrand, sstrand, chromosome, location, probe_set_id, rsid = map(lambda x: index_dict[x], ["snp", "istrand", "sstrand", "chromosome", "location", "probe_set_id", "rsid"])
        af = open(self.annotfile)
        atf = open(annotablefile, 'w')
        discard_lines = self.discard_lines_before_annotate_header();
        for i in range(discard_lines):
            af.readline()
        header = af.readline()
        i = 0
        insert_dict_list = []
        print "reading lines from calls file..."
        discard_chars = string.whitespace + ','
        for line in af:
            i = i+1
            line = line.lstrip(discard_chars).rstrip(discard_chars).split(',')
            if line[0] == "[Controls]":
                break
            fidval = line[probe_set_id].lstrip('"').rstrip('"')
            rsidval = line[rsid].lstrip('"').rstrip('"')
            chromval = line[chromosome].lstrip('"').rstrip('"')
            locationval = line[location].lstrip('"').rstrip('"')
            snpval =  line[snp].lstrip('"').rstrip('"')
            istrandval =  line[istrand].lstrip('"').rstrip('"')
            sstrandval =  line[sstrand].lstrip('"').rstrip('"')
            try:
                #alleleAval, alleleBval = alleles_for_anno_table(snpval, strandval)
                allelea, alleleb = snpval.lstrip("[").rstrip("]").split("/")
            except:
                print line
                raise
            if chromval != '---' and rsidval != '---':
                # make sure both chromosome and location values are integers or sorting won't work.
                allelea_forward, alleleb_forward = self.alleles_trans(allelea, alleleb, istrandval, sstrandval, alleletype="forward")
                allelea_top, alleleb_top = self.alleles_trans(allelea, alleleb, istrandval, sstrandval, alleletype="top")
                d = {'fid': fidval, 'rsid' : rsidval, 'chromosome': self.chromosome_sub(chromval),
                     'location': int(locationval),
                     'allelea_forward_id': allelea_forward, 'alleleb_forward_id': alleleb_forward,
                     'allelea_top_id': allelea_top, 'alleleb_top_id': alleleb_top}
                insert_dict_list.append(d)
        print "done reading lines from calls file, sorting anno table data before insert into database..."
        insert_dict_list.sort(self.sort_anno_dict_list)
        for i, d in enumerate(insert_dict_list):
            atf.write(','.join([str(i+1), d['fid'], d['rsid'], str(d['chromosome']), str(d['location']), d['allelea_forward_id'], d['alleleb_forward_id'], d['allelea_top_id'], d['alleleb_top_id']])+'\n')
        print "finished writing anno table to file..."
        atf.close()

    @print_timing
    def drop_anno_table_constraints_and_indexes(self, session):
        print "dropping anno table constraints..."
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_fid_key', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_pkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_allelea_forward_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_alleleb_forward_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_allelea_top_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_alleleb_top_id_fkey', '%s');"%self.schema)
        session.execute("select * from drop_constraint_if_exists('anno', 'anno_chromosome_fkey', '%s');"%self.schema)
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_allelea_forward_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_alleleb_forward_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_allelea_top_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_alleleb_top_id;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_chromosome;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_location;"%(self.schema, self.schema))
        session.execute("DROP INDEX IF EXISTS %s.ix_%s_anno_rsid;"%(self.schema, self.schema))
        session.commit()
        session.flush()

    @print_timing
    def restore_anno_table_constraints_and_indexes(self, session):
        print "restoring anno table constraints..."
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_fid_key UNIQUE (fid);"%(self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_pkey PRIMARY KEY (id);"%(self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_allelea_forward_id_fkey FOREIGN KEY (allelea_forward_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_alleleb_forward_id_fkey FOREIGN KEY (alleleb_forward_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_allelea_top_id_fkey FOREIGN KEY (allelea_top_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_alleleb_top_id_fkey FOREIGN KEY (alleleb_top_id) REFERENCES %s.allele(val) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("ALTER TABLE ONLY %s.anno ADD CONSTRAINT anno_chromosome_fkey FOREIGN KEY (chromosome) REFERENCES %s.chromo(id) ON UPDATE CASCADE ON DELETE CASCADE;"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_allelea_forward_id ON %s.anno USING btree (allelea_forward_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_alleleb_forward_id ON %s.anno USING btree (alleleb_forward_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_allelea_top_id ON %s.anno USING btree (allelea_top_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_alleleb_top_id ON %s.anno USING btree (alleleb_top_id);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_chromosome ON %s.anno USING btree (chromosome);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_location ON %s.anno USING btree (location);"%(self.schema, self.schema))
        session.execute("CREATE INDEX ix_%s_anno_rsid ON %s.anno USING btree (rsid);"%(self.schema, self.schema))
        session.commit()
        session.flush()
